"""
AI服务路由模块
实现AI相关接口，如数据分析、模型训练等
"""

# 添加项目根目录到Python路径
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 导入FastAPI相关组件
from fastapi import APIRouter, Depends, HTTPException, status, UploadFile, File
# 导入Pydantic模型用于数据验证
from pydantic import BaseModel
# 导入类型提示
from typing import List, Optional, Dict, Any, AsyncGenerator
# 导入JSON处理模块
import json
# 导入pandas用于数据处理
import pandas as pd
# 导入日期时间处理模块
from datetime import datetime
# 导入SQLAlchemy会话
from sqlalchemy.orm import Session

# 导入模型和数据库会话
from utils.models import get_db
# 导入AI客户端
from features.ai.ai_client import ai_client
# 导入配置
from utils.config import Config

# 创建API路由实例
router = APIRouter(
    prefix="/api/v1/ai",  # 路由前缀
    tags=["ai"],  # 标签
    responses={404: {"description": "Not found"}},  # 响应定义
)

# 数据模型定义
class AIRequest(BaseModel):
    """
    AI请求数据模型
    """
    prompt: str            # 提示词
    context: Optional[dict] = None  # 上下文
    model: Optional[str] = None  # 模型名称
    provider: Optional[str] = None  # AI服务提供商
    temperature: Optional[float] = 0.7  # 温度参数
    max_tokens: Optional[int] = None  # 最大令牌数

class AIResponse(BaseModel):
    """
    AI响应数据模型
    """
    message: str          # 响应消息
    response: str         # AI回复
    usage: Optional[dict] = None    # 使用情况
    model: Optional[str] = None     # 使用的模型
    provider: Optional[str] = None  # 使用的提供商

class WorkflowGenerationRequest(BaseModel):
    """
    工作流生成请求数据模型
    """
    description: str        # 描述
    domain: Optional[str] = "bioinformatics"  # 领域
    model: Optional[str] = None  # 模型名称
    provider: Optional[str] = None  # AI服务提供商

class WorkflowNodeTemplate(BaseModel):
    """
    工作流节点模板数据模型
    """
    id: str                 # 节点ID
    type: str               # 节点类型
    position: dict          # 节点位置
    data: dict              # 节点数据

class WorkflowEdgeTemplate(BaseModel):
    """
    工作流边模板数据模型
    """
    id: str                 # 边ID
    source: str             # 源节点ID
    target: str             # 目标节点ID
    type: str               # 边类型

class GeneratedWorkflow(BaseModel):
    """
    生成的工作流数据模型
    """
    name: str                           # 工作流名称
    description: str                    # 工作流描述
    nodes: List[WorkflowNodeTemplate]   # 节点列表
    edges: List[WorkflowEdgeTemplate]   # 边列表

class WorkflowGenerationResponse(BaseModel):
    """
    工作流生成响应数据模型
    """
    message: str                        # 响应消息
    workflow: Optional[GeneratedWorkflow] = None  # 生成的工作流
    usage: Optional[dict] = None        # 使用情况
    model: Optional[str] = None         # 使用的模型
    provider: Optional[str] = None      # 使用的提供商

class ParameterRecommendationRequest(BaseModel):
    """
    参数推荐请求数据模型
    """
    task_type: str          # 任务类型
    data_info: dict         # 数据信息
    previous_results: Optional[dict] = None  # 前期结果
    model: Optional[str] = None  # 模型名称
    provider: Optional[str] = None  # AI服务提供商

class ParameterRecommendationResponse(BaseModel):
    """
    参数推荐响应数据模型
    """
    message: str            # 响应消息
    recommendations: dict   # 推荐参数
    usage: Optional[dict] = None        # 使用情况
    model: Optional[str] = None         # 使用的模型
    provider: Optional[str] = None      # 使用的提供商

class StreamAIRequest(BaseModel):
    """
    流式AI请求数据模型
    """
    prompt: str            # 提示词
    context: Optional[dict] = None  # 上下文
    model: Optional[str] = None  # 模型名称
    provider: Optional[str] = None  # AI服务提供商
    temperature: Optional[float] = 0.7  # 温度参数
    max_tokens: Optional[int] = None  # 最大令牌数

# AI服务类
class AIService:
    """
    AI服务类
    提供AI相关的业务逻辑处理
    """
    
    @staticmethod
    async def chat_with_ai(
        prompt: str, 
        context: Optional[dict] = None, 
        model: Optional[str] = None,
        provider: Optional[str] = None,
        temperature: float = 0.7,
        max_tokens: Optional[int] = None
    ) -> dict:
        """
        与AI进行对话
        Args:
            prompt (str): 用户提示词
            context (Optional[dict]): 对话上下文
            model (Optional[str]): 使用的模型
            provider (Optional[str]): AI服务提供商
            temperature (float): 温度参数
            max_tokens (Optional[int]): 最大令牌数
        Returns:
            dict: AI回复结果
        """
        # 构建消息历史
        messages = []
        
        # 添加上下文信息
        if context:
            if "system_prompt" in context:
                messages.append({"role": "system", "content": context["system_prompt"]})
            if "history" in context:
                messages.extend(context["history"])
        
        # 添加用户当前问题
        messages.append({"role": "user", "content": prompt})
        
        # 调用AI客户端
        result = await ai_client.chat_completion(
            messages=messages,
            model=model,
            provider=provider,
            temperature=temperature,
            max_tokens=max_tokens
        )
        
        return result
    
    @staticmethod
    async def generate_workflow(
        description: str, 
        domain: str = "bioinformatics",
        model: Optional[str] = None,
        provider: Optional[str] = None
    ) -> dict:
        """
        根据自然语言生成工作流
        Args:
            description (str): 工作流描述
            domain (str): 领域
            model (Optional[str]): 使用的模型
            provider (Optional[str]): AI服务提供商
        Returns:
            dict: 生成的工作流
        """
        # 构建系统提示词
        system_prompt = f"""
        你是一个专业的生物信息学工作流设计助手。请根据用户提供的描述生成一个完整的工作流。
        
        领域: {domain}
        
        要求:
        1. 分析用户需求并生成合适的工作流
        2. 工作流应包含合理的节点和连接关系
        3. 节点应包含具体的工具和参数信息
        4. 以JSON格式返回工作流结构
        5. 工作流应包含name、description、nodes和edges字段
        6. nodes中的每个节点应包含id、type、position和data字段
        7. edges中的每个边应包含id、source、target和type字段
        
        示例格式:
        {{
            "name": "工作流名称",
            "description": "工作流描述",
            "nodes": [
                {{
                    "id": "node-1",
                    "type": "nodeType",
                    "position": {{"x": 100, "y": 100}},
                    "data": {{"label": "节点标签", "其他参数": "值"}}
                }}
            ],
            "edges": [
                {{
                    "id": "edge-1",
                    "source": "node-1",
                    "target": "node-2",
                    "type": "smoothstep"
                }}
            ]
        }}
        """
        
        # 构建用户提示词
        user_prompt = f"请为以下需求生成一个{domain}领域的工作流: {description}"
        
        # 构建消息历史
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ]
        
        # 调用AI客户端
        result = await ai_client.chat_completion(
            messages=messages,
            model=model,
            provider=provider,
            temperature=0.5,  # 降低随机性以获得更一致的结构
            max_tokens=2000
        )
        
        # 尝试解析返回的JSON
        try:
            # 提取JSON部分
            response_content = result["response"]
            # 简化处理，实际项目中可能需要更复杂的JSON提取逻辑
            if response_content.startswith("```json"):
                json_str = response_content[7:]
                if json_str.endswith("```"):
                    json_str = json_str[:-3]
            else:
                json_str = response_content
            
            workflow = json.loads(json_str)
            
            # 添加使用情况信息
            workflow["usage"] = result["usage"]
            workflow["model"] = result["model"]
            workflow["provider"] = result["provider"]
            
            return workflow
        except json.JSONDecodeError:
            # 如果解析失败，返回默认工作流
            default_workflow = {
                "name": "默认工作流",
                "description": description,
                "nodes": [
                    {
                        "id": "node-1",
                        "type": "dataInput",
                        "position": {"x": 100, "y": 100},
                        "data": {"label": "数据输入"}
                    },
                    {
                        "id": "node-2",
                        "type": "dataProcess",
                        "position": {"x": 300, "y": 100},
                        "data": {"label": "数据处理"}
                    },
                    {
                        "id": "node-3",
                        "type": "analysis",
                        "position": {"x": 500, "y": 100},
                        "data": {"label": "数据分析"}
                    }
                ],
                "edges": [
                    {"id": "edge-1", "source": "node-1", "target": "node-2", "type": "smoothstep"},
                    {"id": "edge-2", "source": "node-2", "target": "node-3", "type": "smoothstep"}
                ],
                "usage": result["usage"],
                "model": result["model"],
                "provider": result["provider"]
            }
            return default_workflow
    
    @staticmethod
    async def recommend_parameters(
        task_type: str, 
        data_info: dict, 
        previous_results: Optional[dict] = None,
        model: Optional[str] = None,
        provider: Optional[str] = None
    ) -> dict:
        """
        参数推荐
        Args:
            task_type (str): 任务类型
            data_info (dict): 数据信息
            previous_results (Optional[dict]): 前期结果
            model (Optional[str]): 使用的模型
            provider (Optional[str]): AI服务提供商
        Returns:
            dict: 推荐参数
        """
        # 构建系统提示词
        system_prompt = """
        你是一个专业的生物信息学参数推荐助手。请根据任务类型和数据信息提供合适的参数推荐。
        
        要求:
        1. 分析任务类型和数据特征
        2. 提供详细的参数配置建议
        3. 解释每个参数的作用和推荐理由
        4. 以JSON格式返回参数配置
        5. 包含usage、model、provider信息
        
        示例格式:
        {
            "parameters": {
                "param1": {
                    "value": "推荐值",
                    "description": "参数说明",
                    "reason": "推荐理由"
                }
            }
        }
        """
        
        # 构建用户提示词
        user_prompt = f"""
        任务类型: {task_type}
        数据信息: {json.dumps(data_info, ensure_ascii=False, indent=2)}
        前期结果: {json.dumps(previous_results, ensure_ascii=False, indent=2) if previous_results else "无"}
        
        请为上述任务提供参数推荐。
        """
        
        # 构建消息历史
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ]
        
        # 调用AI客户端
        result = await ai_client.chat_completion(
            messages=messages,
            model=model,
            provider=provider,
            temperature=0.3,  # 降低随机性以获得更一致的参数
            max_tokens=1500
        )
        
        # 尝试解析返回的JSON
        try:
            # 提取JSON部分
            response_content = result["response"]
            # 简化处理，实际项目中可能需要更复杂的JSON提取逻辑
            if response_content.startswith("```json"):
                json_str = response_content[7:]
                if json_str.endswith("```"):
                    json_str = json_str[:-3]
            else:
                json_str = response_content
            
            recommendations = json.loads(json_str)
            
            # 添加使用情况信息
            recommendations["usage"] = result["usage"]
            recommendations["model"] = result["model"]
            recommendations["provider"] = result["provider"]
            
            return recommendations
        except json.JSONDecodeError:
            # 如果解析失败，返回默认参数
            default_recommendations = {
                "parameters": {
                    "learning_rate": {
                        "value": 0.001,
                        "description": "学习率",
                        "reason": "适用于大多数深度学习任务的默认值"
                    },
                    "batch_size": {
                        "value": 32,
                        "description": "批处理大小",
                        "reason": "平衡内存使用和训练效率的常用值"
                    }
                },
                "usage": result["usage"],
                "model": result["model"],
                "provider": result["provider"]
            }
            return default_recommendations
    
    @staticmethod
    async def stream_response(
        prompt: str,
        context: Optional[dict] = None,
        model: Optional[str] = None,
        provider: Optional[str] = None,
        temperature: float = 0.7,
        max_tokens: Optional[int] = None
    ) -> AsyncGenerator[str, None]:
        """
        流式输出响应
        Args:
            prompt (str): 用户提示词
            context (Optional[dict]): 对话上下文
            model (Optional[str]): 使用的模型
            provider (Optional[str]): AI服务提供商
            temperature (float): 温度参数
            max_tokens (Optional[int]): 最大令牌数
        Yields:
            str: 流式输出的文本片段
        """
        # 构建消息历史
        messages = []
        
        # 添加上下文信息
        if context:
            if "system_prompt" in context:
                messages.append({"role": "system", "content": context["system_prompt"]})
            if "history" in context:
                messages.extend(context["history"])
        
        # 添加用户当前问题
        messages.append({"role": "user", "content": prompt})
        
        # 调用AI客户端的流式接口
        async for chunk in ai_client.stream_chat_completion(
            messages=messages,
            model=model,
            provider=provider,
            temperature=temperature,
            max_tokens=max_tokens
        ):
            yield chunk
        
        # 结束标记
        yield "\n[END_OF_RESPONSE]"

# 路由定义
@router.post("/chat", response_model=AIResponse)
async def chat_with_ai(request: AIRequest):
    """
    与AI进行对话接口
    集成大模型API进行对话
    Args:
        request (AIRequest): AI请求
    Returns:
        AIResponse: AI响应
    """
    try:
        # 调用AI服务
        result = await AIService.chat_with_ai(
            prompt=request.prompt,
            context=request.context,
            model=request.model,
            provider=request.provider,
            temperature=request.temperature,
            max_tokens=request.max_tokens
        )
        
        return {
            "message": "AI回复成功",
            "response": result["response"],
            "usage": result["usage"],
            "model": result["model"],
            "provider": result["provider"]
        }
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"AI服务调用失败: {str(e)}"
        )

@router.post("/generate-workflow", response_model=WorkflowGenerationResponse)
async def generate_workflow(request: WorkflowGenerationRequest):
    """
    根据自然语言生成工作流接口
    实现自然语言转工作流逻辑
    Args:
        request (WorkflowGenerationRequest): 工作流生成请求
    Returns:
        WorkflowGenerationResponse: 工作流生成响应
    """
    try:
        # 调用AI服务生成工作流
        workflow_result = await AIService.generate_workflow(
            description=request.description,
            domain=request.domain,
            model=request.model,
            provider=request.provider
        )
        
        # 提取工作流数据
        workflow_data = {
            "name": workflow_result.get("name", "生成的工作流"),
            "description": workflow_result.get("description", request.description),
            "nodes": [WorkflowNodeTemplate(**node) for node in workflow_result.get("nodes", [])],
            "edges": [WorkflowEdgeTemplate(**edge) for edge in workflow_result.get("edges", [])]
        }
        
        return {
            "message": "工作流生成成功",
            "workflow": GeneratedWorkflow(**workflow_data),
            "usage": workflow_result.get("usage"),
            "model": workflow_result.get("model"),
            "provider": workflow_result.get("provider")
        }
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"工作流生成失败: {str(e)}"
        )

@router.post("/recommend-parameters", response_model=ParameterRecommendationResponse)
async def recommend_parameters(request: ParameterRecommendationRequest):
    """
    参数推荐接口
    实现参数推荐逻辑
    Args:
        request (ParameterRecommendationRequest): 参数推荐请求
    Returns:
        ParameterRecommendationResponse: 参数推荐响应
    """
    try:
        # 调用AI服务推荐参数
        recommendations_result = await AIService.recommend_parameters(
            task_type=request.task_type,
            data_info=request.data_info,
            previous_results=request.previous_results,
            model=request.model,
            provider=request.provider
        )
        
        return {
            "message": "参数推荐完成",
            "recommendations": recommendations_result.get("parameters", {}),
            "usage": recommendations_result.get("usage"),
            "model": recommendations_result.get("model"),
            "provider": recommendations_result.get("provider")
        }
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"参数推荐失败: {str(e)}"
        )

@router.post("/stream")
async def stream_response(request: StreamAIRequest):
    """
    流式输出接口
    实现流式输出功能（Server-Sent Events）
    Args:
        request (StreamAIRequest): 流式AI请求
    Returns:
        StreamingResponse: 流式响应
    """
    try:
        # 创建流式响应
        return StreamingResponse(
            AIService.stream_response(
                prompt=request.prompt,
                context=request.context,
                model=request.model,
                provider=request.provider,
                temperature=request.temperature,
                max_tokens=request.max_tokens
            ),
            media_type="text/plain"
        )
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"流式输出失败: {str(e)}"
        )